import os
import shutil
from xml.etree import ElementTree as ET
from xml.etree.ElementTree import Element, SubElement
from xml.dom import minidom
import cv2


def write_xml(xml_name):
    # tree = ET.parse(xml_name)
    # root = tree.getroot()
    # filename = root.find('Object').text
    # print(filename)

    tree = ET.parse(xml_name)
    root = tree.getroot()
    # imagename = root.find('Object')
    #
    #
    # size = root.find('size')
    # w = int(size.find('width').text)
    # h = int(size.find('height').text)
    index_num = 0
    for obj in root.iter('Object'):
        cls = obj.find('Rect').text
        index_num += 1
        print(index_num)
        print(obj)


def usc_xml_read(xml_name):
    with open(xml_name, 'r') as xml_file:
        lines = xml_file.readlines()
        boxes = []
        for line in lines:
            if 'Rect' in line:
                line = line.strip()[1:-2]
                print(line.strip())
                texts = line.split(" ")
                x, y, width, height = 0, 0, 0, 0
                for text in texts:
                    if not ("x" in text or "y" in text or "width" in text or "height" in text):
                        continue
                    if "x" in text:
                        x = text.strip().split("=")[-1][1:-1]
                    if "y" in text:
                        y = text.strip().split("=")[-1][1:-1]
                    if "width" in text:
                        width = text.strip().split("=")[-1][1:-1]
                    if "height" in text:
                        height = text.strip().split("=")[-1][1:-1]
                print("x:", x, "y:", y, "width:", width, "height:", height)
                boxes.append(['person', x, y, width, height])
        return boxes


def get_xml(img_name,size,roi,outpath):
    '''
    传入每张图片的名称 大小 和标记点列表生成和图片名相同的xml标记文件
    :param img_name: 图片名称
    :param size:[width,height,depth]
    :param roi:[["label1",xmin,ymin,xmax,ymax]["label2",xmin,ymin,xmax,ymax].....]
    :return:
    '''
    impl = minidom.getDOMImplementation()
    doc = impl.createDocument(None, None, None)
    #创建根节点
    orderlist = doc.createElement("annotation")
    doc.appendChild(orderlist)
    #c创建二级节点
    filename=doc.createElement("filename")
    filename.appendChild(doc.createTextNode(img_name))
    orderlist.appendChild(filename)
    #size节点
    sizes=doc.createElement("size")
    width=doc.createElement("width")
    width.appendChild(doc.createTextNode(str(size[0])))
    height=doc.createElement("height")
    height.appendChild(doc.createTextNode(str(size[1])))
    depth=doc.createElement("depth")
    depth.appendChild(doc.createTextNode(str(size[2])))
    sizes.appendChild(width)
    sizes.appendChild(height)
    sizes.appendChild(depth)
    orderlist.appendChild(sizes)
    #object 节点
    for ri in roi:
        object=doc.createElement("object")
        name=doc.createElement("name")
        name.appendChild(doc.createTextNode(ri[0]))
        object.appendChild(name)
        bndbox=doc.createElement("bndbox")

        xmin=doc.createElement("xmin")
        xmin.appendChild(doc.createTextNode(str(ri[1])))
        bndbox.appendChild(xmin)
        ymin=doc.createElement("ymin")
        ymin.appendChild(doc.createTextNode(str(ri[2])))
        bndbox.appendChild(ymin)
        xmax=doc.createElement("xmax")
        xmax.appendChild(doc.createTextNode(str(ri[3])))
        bndbox.appendChild(xmax)
        ymax=doc.createElement("ymax")
        ymax.appendChild(doc.createTextNode(str(ri[4])))
        bndbox.appendChild(ymax)
        object.appendChild(bndbox)
        orderlist.appendChild(object)
         # 将dom对象写入本地xml文件

    # 打开test.xml文件 准备写入
    f = open(os.path.join(outpath,img_name[:-4]+'.xml'), 'w')
    # 写入文件
    doc.writexml(f, addindent='  ', newl='\n')
    # 关闭
    f.close()


if __name__ == '__main__':
    src_xml_path = 'D:/BaiduNetdiskDownload/USC行人数据集/dataA/xml_src/9Cup10.gt.xml'
    dst_xml_path = 'D:/BaiduNetdiskDownload/USC行人数据集/dataA/xml_dst_1/'
    dst_xml_name = '9Cup10.xml'

    images_path = "D:/BaiduNetdiskDownload/USC行人数据集/dataA/images"
    image_name = dst_xml_name[:-4]+'.bmp'
    image_file_path = os.path.join(images_path, image_name)
    image = cv2.imread(image_file_path)
    height, width, channel = image.shape
    # get_usc_xml(xml_path)
    boxes = usc_xml_read(src_xml_path)
    get_xml(dst_xml_name, [height, width, channel], boxes, dst_xml_path)
    print(boxes)